# Copyright 2024 Marimo. All rights reserved.
from __future__ import annotations

import abc
import base64
import copy
import sys
import types
import uuid
import weakref
from dataclasses import dataclass, fields
from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    Generic,
    Optional,
    Sequence,
    TypeVar,
    cast,
)

from marimo import _loggers
from marimo._output.hypertext import Html
from marimo._output.rich_help import mddoc
from marimo._plugins.core.web_component import JSONType, build_ui_plugin
from marimo._plugins.ui._core import ids
from marimo._runtime.context import (
    ContextNotInitializedError,
    RuntimeContext,
    get_context,
)
from marimo._runtime.functions import Function

if TYPE_CHECKING:
    from marimo._plugins.ui._impl.input import form as form_plugin

# S: Type of frontend value
#   - the initial value sent to the frontend must be of type S
#   - values received from the frontend must be of type S
S = TypeVar("S", bound=JSONType)

# T: Type of `UIElement.value`
#   - UIElement's `_convert_value` method is used to convert a frontend
#   value (of type S) to a value of type T
T = TypeVar("T")

LOGGER = _loggers.marimo_logger()


@dataclass
class Lens:
    """Track how a view of a higher-order element relates to its source

    Higher-order UI elements support lensing, ie extracting their children
    as "views". These views can be embedded in other outputs and interacted
    with.

    UI elements that are views of a higher-order element (eg, an entry of
    an array is a view of the array) have a lens object that stores the
    id of its parent UI element, and the key at which its parent stores it.
    """

    parent_id: str
    key: str


@dataclass
class InitializationArgs(Generic[S, T]):
    component_name: str
    initial_value: S
    label: Optional[str]
    on_change: Optional[Callable[[T], None]]
    args: dict[str, JSONType]
    slotted_html: str
    functions: tuple[Function[Any, Any], ...]


class MarimoConvertValueException(Exception):
    pass


@mddoc
class UIElement(Html, Generic[S, T], metaclass=abc.ABCMeta):
    """An HTML element with a value

    A `UIElement` is an HTML element with a value; when the value of the
    element on the page changes, the value of the UIElement is updated as well.

    This is an abstract class. `UIElement`s are responsible for mapping values
    sent by the frontend (of type S) to values expected by the Python object
    (of type T): a subclass can be made concrete by implementing the
    `_convert_value` method.

    Type Parameters:

    - S: The type of the values sent by the frontend to the kernel; must be
         JSON-serializable
    - T: The type of the UIElement's value; can be any type

    **Attributes.**

    - value: The value of the `UIElement`.

    **Methods.**

    - form: create a submittable form this `UIElement`.
    """

    _value_frontend: S
    _value: T

    def __init__(
        self,
        component_name: str,
        initial_value: S,
        label: Optional[str],
        on_change: Optional[Callable[[T], None]],
        args: dict[str, JSONType],
        slotted_html: str = "",
        functions: tuple[Function[Any, Any], ...] = (),
    ) -> None:
        """Initialize a UIElement

        Args:
        ----
        component_name: tag name of the custom element
        initial_value: initial value of the element in the frontend
        label: markdown string, label of element
        on_change: callback, called with element's new value on change
        args: arguments that the element takes
        slotted_html: any html to slot in the custom element
        functions: any functions to register with the graph
        """
        # Validate parameters from a user
        if not isinstance(component_name, str):
            raise TypeError("component_name must be a string")
        if label is not None and not isinstance(label, str):
            raise TypeError("label must be a string or None")
        if on_change is not None and not callable(on_change):
            raise TypeError("on_change must be a callable or None")

        # arguments stored in signature order for cloning
        self._component_args = args
        self._args: InitializationArgs[S, T] = InitializationArgs(
            component_name=component_name,
            initial_value=initial_value,
            label=label,
            on_change=on_change,
            args=args,
            slotted_html=slotted_html,
            functions=functions,
        )
        self._initialized = False
        self._initialize(self._args)
        self._initialized = True

    def _initialize(
        self, initialization_args: InitializationArgs[S, T]
    ) -> None:
        """Initialize the UIElement

        Split out from __init__ so _clone() typechecks
        """
        (
            component_name,
            initial_value,
            label,
            on_change,
            args,
            slotted_html,
            functions,
        ) = (
            initialization_args.component_name,
            initialization_args.initial_value,
            initialization_args.label,
            initialization_args.on_change,
            initialization_args.args,
            initialization_args.slotted_html,
            initialization_args.functions,
        )
        # A UIElement may be a child ("lens") of another UI element.
        #
        # Set with self._register_as_view() after initialization, since parents
        # are usually created after the child is created
        self._lens: Lens | None = None

        # Random token
        #
        # Every element is annotated with a random token, which by design is
        # different every time the element is constructed (i.e., every time a
        # cell runs): this guarantees that re-running a cell that creates a UI
        # element will trigger a re-render and reset it to its initial value.
        # We need this to ensure that the element on the page is synchronized
        # with the element in the kernel.
        self._random_id = str(uuid.uuid4())

        # Stable ID
        #
        # Every element has an ID that is used for two purposes:
        #  - to synchronize multiple instances of the element on the page
        #  - to synchronize elements on the page with elements in the kernel
        #
        # IDs are stable across multiple sessions if the set of UI elements
        # created by each cell is deterministic; this fact is used to
        # optionally override the element's initial value.
        try:
            self._id = get_context().take_id()
        except (ids.NoIDProviderException, ContextNotInitializedError):
            self._id = self._random_id

        self._ctx: RuntimeContext | None
        try:
            # cache the context in case the UI element is constructed
            # in a nested context -- so that if the UI element is accessed
            # in the root context (eg with app_result.defs["elem"].value),
            # the correct constructing context is retrieved
            self._ctx = get_context()
        except ContextNotInitializedError:
            self._ctx = None
        else:
            # When the UI element is destructed, it should be removed
            # from the UIElementRegistry (which only holds a weakref to it).
            finalizer = weakref.finalize(
                self, self._ctx.ui_element_registry.delete, self._id, id(self)
            )
            # No need to clean up the registry at program teardown
            finalizer.atexit = False

        if self._ctx is not None:
            self._ctx.ui_element_registry.register(self._id, self)
            # an Instantiate request may want us to override the initial value
            try:
                # NB: If a cell produces a non-deterministic set of
                # UI elements, a UI element may be matched with an initial
                # value that was actually for some other element
                #
                # TODO(akshayka): validate the tag-name to make sure that the
                # value is at least the right type (ie, S)
                #
                # TODO(akshayka): parametrize UIElement with an optional
                # string ID, so users can provide their own IDs to make
                # sure a mismatch never happens ...
                initial_value = cast(
                    S, self._ctx.get_ui_initial_value(self._id)
                )
            except KeyError:
                # we weren't asked to override the UI element's value
                pass

            for function in functions:
                self._ctx.function_registry.register(
                    namespace=self._id, function=function
                )
        self._initial_value_frontend = initial_value
        self._value_frontend = initial_value
        self._value = self._initial_value = self._convert_value(initial_value)
        self._on_change = on_change
        self._component_args = args

        self._inner_text = build_ui_plugin(
            component_name,
            initial_value,
            label,
            args,
            slotted_html,
        )
        text = (
            f"<marimo-ui-element object-id='{self._id}' "
            + f"random-id='{self._random_id}'>"
            + self._inner_text
            + "</marimo-ui-element>"
        )
        super().__init__(text=text)

    @abc.abstractmethod
    def _convert_value(self, value: S) -> T:
        """Converts a value from the frontend to a value for the `UIElement`

        This method must convert `value`, the JSON-decoded value sent by the
        frontend, to a value of type `T` for the `UIElement`.
        """
        pass

    def _register_as_view(self, parent: UIElement[Any, Any], key: str) -> None:
        """Register this element as a view of `parent`."""
        self._lens = Lens(parent_id=parent._id, key=key)

    @property
    def value(self) -> T:
        """The element's current value."""
        if self._ctx is None:
            return self._value

        if (
            self._ctx.execution_context is not None
            and not self._ctx.execution_context.setting_element_value
            and (
                self._ctx.execution_context.cell_id
                == self._ctx.ui_element_registry.get_cell(self._id)
            )
        ):
            raise RuntimeError(
                "Accessing the value of a UIElement in the cell that created "
                "it is not allowed. Fix: move the value access to another "
                "cell."
            )
        return self._value

    @value.setter
    def value(self, value: T) -> None:
        del value
        raise RuntimeError(
            "Setting the value of a UIElement is not allowed. "
            "If you need to imperatively set the value of a UIElement, "
            "consider using mo.state()."
        )

    def __setattr__(self, name: str, value: Any) -> None:
        if name == "on_change":
            raise RuntimeError(
                "Setting the on_change handler of a UIElement is not allowed. "
                "You must set the on_change in the constructor."
            )
        super().__setattr__(name, value)

    @mddoc
    def form(
        self,
        label: str = "",
        *,
        bordered: bool = True,
        loading: bool = False,
        submit_button_label: str = "Submit",
        submit_button_tooltip: Optional[str] = None,
        submit_button_disabled: bool = False,
        clear_on_submit: bool = False,
        show_clear_button: bool = False,
        clear_button_label: str = "Clear",
        clear_button_tooltip: Optional[str] = None,
        validate: Optional[
            Callable[[Optional[JSONType]], Optional[str]]
        ] = None,
        on_change: Optional[Callable[[Optional[T]], None]] = None,
    ) -> form_plugin[S, T]:
        """Create a submittable form out of this `UIElement`.

        Use this method to create a form that gates the submission
        of a `UIElement`s value until a submit button is clicked.

        The value of the `form` is the value of the underlying
        element the last time the form was submitted.

        **Examples.**

        Convert any `UIElement` into a form:

        ```python
        prompt = mo.ui.text_area().form()
        ```

        Combine with `HTML.batch` to create a form made out of multiple
        `UIElements`:

        ```python
        form = (
            mo.ui.md(
                '''
            **Enter your prompt.**

            {prompt}

            **Choose a random seed.**

            {seed}
            '''
            )
            .batch(
                prompt=mo.ui.text_area(),
                seed=mo.ui.number(),
            )
            .form()
        )
        ```

        **Args.**

        - `label`: A text label for the form.
        - `bordered`: whether the form should have a border
        - `loading`: whether the form should be in a loading state
        - `submit_button_label`: the label of the submit button
        - `submit_button_tooltip`: the tooltip of the submit button
        - `submit_button_disabled`: whether the submit button should be
          disabled
        - `clear_on_submit`: whether the form should clear its contents after
            submitting
        - `show_clear_button`: whether the form should show a clear button
        - `clear_button_label`: the label of the clear button
        - `clear_button_tooltip`: the tooltip of the clear button
        - `validate`: a function that takes the form's value and returns an
            error message if the value is invalid,
            or `None` if the value is valid
        """
        from marimo._plugins.ui._impl.input import form as form_plugin

        return form_plugin(
            element=self,
            label=label,
            bordered=bordered,
            loading=loading,
            submit_button_label=submit_button_label,
            submit_button_tooltip=submit_button_tooltip,
            submit_button_disabled=submit_button_disabled,
            clear_on_submit=clear_on_submit,
            show_clear_button=show_clear_button,
            clear_button_label=clear_button_label,
            clear_button_tooltip=clear_button_tooltip,
            validate=validate,
            on_change=on_change,
        )

    def send_message(
        self, message: JSONType, buffers: Optional[Sequence[bytes]]
    ) -> None:
        """
        Send a message to the element rendered on the frontend
        from the backend.
        """

        from marimo._messaging.ops import SendUIElementMessage

        SendUIElementMessage(
            ui_element=self._id,
            message=message,
            buffers=[
                base64.b64encode(buffer).decode() for buffer in (buffers or [])
            ],
        ).broadcast()

    def _update(self, value: S) -> None:
        """Update value, given a value from the frontend

        Calls the on_change handler with the element's new value as a
        side-effect.
        """
        self._value_frontend = value
        try:
            self._value = self._convert_value(value)
        except MarimoConvertValueException:
            raise

        if self._on_change is not None:
            self._on_change(self._value)

    def _on_update_completion(self) -> None:
        """Callback to run after the kernel has processed a value update."""
        return

    def __deepcopy__(self, memo: dict[int, Any]) -> UIElement[S, T]:
        # Custom deepcopy that excludes elements that can't be deepcopied
        cls = self.__class__
        result = cls.__new__(cls)
        memo[id(self)] = result
        for k, v in self.__dict__.items():
            if isinstance(v, RuntimeContext):
                setattr(result, k, v)
            else:
                setattr(result, k, copy.deepcopy(v, memo))

        # Get a new object ID and function namespace
        #
        # We use the new instance's functions, since they are typically bound
        # to the UI element instance. But we only use the new on_change
        # if the old one was bound to self.
        args: InitializationArgs[S, T]
        if (
            isinstance(self._args.on_change, types.MethodType)
            and self._args.on_change.__self__ is self
        ):
            # on_change was bound to self; use the new one.
            args = InitializationArgs(
                **{
                    # dataclass asdict does a deepcopy, we want shallow.
                    **{
                        field.name: getattr(self._args, field.name)
                        for field in fields(self._args)
                    },
                    "on_change": result._args.on_change,
                    "functions": result._args.functions,
                }
            )
        else:
            # otherwise, use the original on_change, which may be a state
            # SetFunctor or something else unrelated to this instance.
            args = InitializationArgs(
                **{
                    **{
                        field.name: getattr(self._args, field.name)
                        for field in fields(self._args)
                    },
                    "functions": result._args.functions,
                }
            )
        result._initialize(args)
        return result

    def _clone(self) -> UIElement[S, T]:
        """Clone a UIElement, returning one with a different id

        The clone will not synchronize with the original element.

        Composite UIElement may need to override this method to run
        their own side-effects.
        """
        return copy.deepcopy(self)

    def __bool__(self) -> bool:
        sys.stderr.write(
            "The truth value of a UIElement is always True. You "
            "probably want to call `.value` instead."
        )
        return True
